import datetime
import os.path
import signal
import sys
import stock_info as si
from multiprocessing import Process, Queue

import pandas as pd

def download_stock(tic, path, days_from_now, append_data=False):
    if os.path.exists(path):
        if append_data:
            df_old = pd.read_csv(path, parse_dates=True, index_col=0)
            start_date = df_old.index[-1].date() + datetime.timedelta(days=1)
            end_date = datetime.date.today()
            if start_date==end_date:
                print('alreay latest')
            else:
                df_new = si.get_data(tic, start_date, end_date)
                df_new = pd.concat([df_old, df_new])
                df_new.to_csv(path, index=True)
                print(f'Append {tic}')
        else:
            print(f'{tic} already exists')
    else:
        start_date = datetime.datetime.now() - datetime.timedelta(days=days_from_now)
        end_date = datetime.datetime.now()
        df_new = si.get_data(tic, start_date, end_date)
        df_new.to_csv(path, index=True)
        print(f'Download {tic}')
        
        

def worker(task_queue, save_root, days_from_now, append_data):
    while not task_queue.empty():
        tic = task_queue.get()
        save_path=os.path.join(save_root, tic + '.csv')
        try:
            download_stock(tic, save_path, days_from_now, append_data)
        except Exception as e:
            print(f'Error: {tic} {e}')
        


def multi_download_stock(tickers, save_root, days_from_now, append_data, num_workers=os.cpu_count()):
    
    q = Queue()
    for tic in sorted(tickers):
        q.put(tic)

    input_dict=dict(
        task_queue=q,
        save_root=save_root,
        days_from_now=days_from_now,
        append_data=append_data
    )
    
    # worker(**input_dict)
    processes = [Process(target=worker, kwargs=input_dict) for _ in range(num_workers)]
    for p in processes:
        p.start()

    def stop_processes(signal, frame):
        print('Received signal, stopping processes...')
        # sys.exit(0)
        for p in processes:
            p.terminate()
            p.join()

        sys.exit(0)
    signal.signal(signal.SIGINT, stop_processes)
    signal.signal(signal.SIGTERM, stop_processes)

    for p in processes:
        p.join()  # 等待每个子进程完成
        
if __name__ == '__main__':
    with open('./all_sym.txt', 'r')as f:
        tickers = f.read().split()
        print(len(tickers))
    
    # tickers=['TSLA',"AAPL"]
    # df=pd.read_csv('../csv/TSLA.csv',parse_dates=True,index_col=0)
    # print(df.head())
    multi_download_stock(tickers,
                         save_root='../csv/',
                         days_from_now=365*10,
                         append_data=True,
                         num_workers=os.cpu_count()
                         )